import os.path
import time

import numpy as np
from docx.document import Document as _Document
from docx.oxml.table import CT_Tbl
from docx.oxml.text.paragraph import CT_P
from docx.table import _Cell, Table
from docx.text.paragraph import Paragraph
from docx import Document
import re
import pandas as pd


# 从block中提取图片文件
def get_imgs(figure_block, doc, caption=''):
    """
    :param figure_block: 待测block
    :param doc: Document对象
    :param caption: 图片名，不指定则以时间戳命名
    :return:
    """
    if caption is None or len(caption) == 0:
        caption = str(int(time.time()))
    imgs = []
    for img in figure_block._element.xpath('.//pic:pic'):
        embed = img.xpath('.//a:blip/@r:embed')[0]
        related_part = doc.part.related_parts[embed]
        suffix = related_part.filename.split('.')[-1]
        if suffix in ('jpg', 'jpeg', 'png', 'jfif', 'exif', 'gif', 'tiff', 'bmp'):
            image = related_part.image
            ext = image.ext
            imgs.append((image, ext))
        else:
            print(f'不能提取图片{caption}，原因：{suffix}文件')
    return imgs, caption


# 处理表格，以字典形式输出表格
def process_table(table):
    table_content = []  # 表数据
    table_heads = []  # 表头
    correct_head_ids = []  # 满足要求的表头id
    for i, row in enumerate(table.rows):
        # 首先获取当前表表头，记录合格的表头id,合格表头不能是'',None
        if i == 0:
            for head_id, head in enumerate(row.cells):
                table_heads.append(head.text)
                if head.text:
                    correct_head_ids.append(head_id)
                    # 用表头封装当前表格内容
        else:
            row_dic = {}
            for cell_id, cell in enumerate(row.cells):
                value = cell.text  # value对应的是每个单元格内容
                # 内容有合格的表头，并且内容本身也合格，合格的内容不能是'',None
                if cell_id in correct_head_ids and value:
                    row_dic[table_heads[cell_id]] = value
            # 行数据不为空时记录改行
            if row_dic:
                table_content.append(row_dic)
    if len(table_content) == 0:
        return None
    else:
        return str(table_content)


# 遍历docx文件中所有的子元素
def iter_block_items(parent):
    if isinstance(parent, _Document):
        parent_elm = parent.element.body
    elif isinstance(parent, _Cell):
        parent_elm = parent._tc
    else:
        raise ValueError("something's not right")

    for child in parent_elm.iterchildren():
        if isinstance(child, CT_P):
            yield Paragraph(child, parent)
        elif isinstance(child, CT_Tbl):
            yield Table(child, parent)


# 文件无格式，直接输出原文
def get_content(doc, img_path):
    """
    :param doc: Document对象
    :param img_path: 图片保存路径
    :return: 纯文本
    """
    content = []
    for block_id, block in enumerate(iter_block_items(doc)):
        if isinstance(block, Paragraph):
            if len(block.text) > 0:
                content.append(block.text)
            # 图片处理，若有图片则将图片保存到img_path
            imgs, caption = get_imgs(block, doc)
            if len(imgs) > 0:
                if not os.path.exists(img_path):
                    os.makedirs(img_path)
                for img, img_suffix in imgs:
                    with open(f'{img_path}/{caption}.{img_suffix}', 'wb') as f:
                        f.write(img.blob)
                    content.append(f'<image:{img_path}/{caption}.{img_suffix}>')

        elif isinstance(block, Table):
            table = process_table(block)
            if table is not None:
                content.append(f'<table:{table}>')
    return '\n'.join(content)


# 删除目录（若有），并转换为dataframe格式
def check_catalog_to_df(result):
    idns = result['章节索引']
    for i, idn in enumerate(idns):
        if idn in idns[1 + i:]:
            result['章节索引'][i] = np.NAN
    df = pd.DataFrame(result)
    df.dropna(inplace=True)
    return df


# 利用正则表达式处理文本，检索
def re_parse(name: str, doc: Document, pattern, img_path):
    """
    :param name: 文件名
    :param doc: Document对象
    :param pattern: 正则表达式
    :return: Dataframe对象
    """
    result = {'文件名': [], '章节索引': [], '标题': [], 'zhengwen': [], '图片路径': [], '图片数量': [], '表格数量': []}
    content = get_content(doc, f'{img_path}/{name}_img')
    section_id = re.findall(pattern, content)
    for i in range(len(section_id)):
        if i < len(section_id) - 1:
            s = content.find(section_id[i])
            e = content.find(section_id[i + 1])
            text = content[s:e].replace(section_id[i], '').strip()
            content = content[e:]
        else:
            text = content.replace(section_id[i], '').strip()
        text_split = text.split('\n')
        if len(text_split) > 0:
            result['标题'].append(text_split[0])
            text = re.sub(pattern, '', text)
            title_content = text.replace(text_split[0], '').strip()
            result['zhengwen'].append(title_content)
        else:
            result['标题'].append('')
            text = re.sub(pattern, '', text)
            title_content = text.strip()
            result['zhengwen'].append(title_content)
        imgs = re.findall('<image:([^>]+)>', title_content)
        tables = re.findall('<table:[^>]+>', title_content)
        result['图片数量'].append(len(imgs))
        result['图片路径'].append('_'.join(imgs))
        result['表格数量'].append(len(tables))
        result['章节索引'].append(section_id[i].strip())
    result['文件名'] = [name + '.docx'] * len(section_id)
    return result


# 返回当前列表中指定元素的最后一个的位置
def get_last_pos(ls, e):
    res = -1
    for pos, i in enumerate(ls):
        if i == e:
            res = pos
    return res


# 根据标题等级重新定义标题，要求子标题要携带所有父标题信息
def redefine_with_level(titles, titles_lvs, split_sign='___'):
    """
    :param titles: 子标题
    :param titles_lvs: 子标题等级
    :param split_sign: 分割符
    :return:
    """
    label_titles = []  # 重新定义的标题
    for id, title in enumerate(titles):
        sub_lv_ls = titles_lvs[0:id]
        cur_lv = titles_lvs[id]
        cur_titles = [title]
        # 子标题找到其最近父标题
        while cur_lv > 1:
            cur_lv = cur_lv - 1
            last_pos = get_last_pos(sub_lv_ls, cur_lv)
            cur_titles.insert(0, titles[last_pos])
        label_titles.append(split_sign.join(cur_titles))
    return label_titles
